{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 10.5 多头注意力"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "from torch import nn\n",
    "from d2l import torch as d2l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 两个重要的转置函数\n",
    "def transpose_qkv(X, num_heads):\n",
    "    # X input:       (batch_size, query_size or kv_size, num_hiddens)\n",
    "    # X after shape: (batch_size, query_size or kv_size, num_heads, num_hiddens / num_heads)\n",
    "    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)\n",
    "    \n",
    "    # X after permute:(batch_size, num_heads, query_size or kv_size, num_hiddens / num_heads)\n",
    "    X = X.permute(0, 2, 1, 3)\n",
    "    \n",
    "    # return output: (batch_size * num_heads, query_size or kv_size, num_hiddens / num_heads)\n",
    "    return X.reshape(-1, X.shape[2], X.shape[3])\n",
    "\n",
    "def transpose_output(X, num_heads):\n",
    "    \"\"\"函数transpose_qkv的逆操作\"\"\"\n",
    "    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])\n",
    "    X = X.permute(0, 2, 1, 3)\n",
    "    return X.reshape(X.shape[0], X.shape[1], -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 多头注意力\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, key_size, query_size, value_size, num_hiddens,\n",
    "                num_heads, dropout, bias=False, **kwargs):\n",
    "        super(MultiHeadAttention, self).__init__(**kwargs)\n",
    "        self.num_heads = num_heads\n",
    "        self.attention = d2l.DotProductAttention(dropout)\n",
    "        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)\n",
    "        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)\n",
    "        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)\n",
    "        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)\n",
    "        \n",
    "    def forward(self, queries, keys, values, valid_lens):\n",
    "        # q,k,v:      (batch_size, qkv_size, num_hiddens)\n",
    "        # valid_lens: (batch_size,) or (batch_size, query_size)\n",
    "        queries = transpose_qkv(self.W_q(queries), self.num_heads)\n",
    "        keys = transpose_qkv(self.W_k(keys), self.num_heads)\n",
    "        values = transpose_qkv(self.W_v(values), self.num_heads)\n",
    "        \n",
    "        if valid_lens is not None:\n",
    "            # 在第0轴，将valid_lens(标量或矢量)复制num_heads次\n",
    "            valid_lens = torch.repeat_interleave(valid_lens, repeats=num_heads, dim=0)\n",
    "            \n",
    "        # output: (batch_size * num_heads, qkv_size, num_hiddens / num_heads)\n",
    "        output = self.attention(queries, keys, values, valid_lens)\n",
    "        \n",
    "        output_concat = transpose_output(output, self.num_heads)\n",
    "        return self.W_o(output_concat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MultiHeadAttention(\n",
       "  (attention): DotProductAttention(\n",
       "    (dropout): Dropout(p=0.5, inplace=False)\n",
       "  )\n",
       "  (W_q): Linear(in_features=128, out_features=100, bias=False)\n",
       "  (W_k): Linear(in_features=128, out_features=100, bias=False)\n",
       "  (W_v): Linear(in_features=128, out_features=100, bias=False)\n",
       "  (W_o): Linear(in_features=100, out_features=100, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_qkv_size, num_hiddens, num_heads = 128, 100, 5\n",
    "attention = MultiHeadAttention(num_qkv_size, num_qkv_size, num_qkv_size, \n",
    "                               num_hiddens, num_heads, 0.5)\n",
    "attention.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 100])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_hiddens, num_heads = 100, 5\n",
    "attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens, \n",
    "                               num_hiddens, num_heads, 0.5)\n",
    "attention.eval()\n",
    "\n",
    "batch_size, num_queries, num_kvpairs, valid_lens = 2, 4, 6, torch.tensor([3, 2])\n",
    "X = torch.ones((batch_size, num_queries, num_hiddens))\n",
    "Y = torch.ones((batch_size, num_kvpairs, num_hiddens))\n",
    "attention(X, Y, Y, valid_lens).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
